/*******************************************************************
 *  File Name: LZ.cpp
 *  Description: Main file of the LZ encode and decode procedure
 *
 *		This project is composed of the following 3 files:
 *		- LZ.cpp			this file, the logic frame
 *		- MainLogic.h		the function bodys of LZ.cpp
 *		- structures.h		the main data structures definitions
 *
 *
 *  Auther: Zhipeng Xu
 *  Last Modified: 2019-11-16 20:43:58
 ******************************************************************/

#include "structures.h"
#include "MainLogic.h"

int main(int argc, char const* argv[])
{
	/* Step 0 - Print the Welcome banner and prompt for filename */
	std::string fname = banner();
	std::string encodedfname = fname + ".encoded.lz";
	std::string decodedfname = fname + ".decoded.lz";
	const char* filename = fname.c_str();
	const char* encodedfilename = encodedfname.c_str();
	const char* decodedfilename = decodedfname.c_str();

	/* Step 1 - read in the binay file to be encoded */
	struct stat filestat;
	stat(filename, &filestat);
	size_t filelength = filestat.st_size;
	char* filecontent = (char*)malloc(filelength * sizeof(char) + 1);

	FILE* fp;
	errno_t err = fopen_s(&fp, filename, "rb");
	if (err) {
		printf("File Does Not Exist!\n");
		exit(-1);
	}
	fread(filecontent, sizeof(char), filelength, fp);
	fclose(fp);
	printf("Successfully Opened %s, size: %lu Bytes\n", filename, (long unsigned)filelength);


	/* Step 2 - prepare dictionary related structres */
	size_t cidx = 1, i = 0;	/* cidx is the current maximum index */
	Entry* dict = (Entry*)malloc(sizeof(Entry) * filelength);
	if (!dict) {
		printf("File too Large!\n");
		printf("Out of memory due to the filelength.\n");
		exit(-1);
	}
	unsigned char* wordbuf = (unsigned char*)malloc(sizeof(unsigned char) * filelength);
	unsigned char* wordbufptr = wordbuf;

	/* use a search tree to speed up the search ancestor procedure */
	TreeNode treehead, * searchptr;
	initnode(&treehead, 0, 0);

	LARGE_INTEGER t1, t2, tc;
	QueryPerformanceFrequency(&tc);
	QueryPerformanceCounter(&t1);
	/* Step 3 - Start making the dictionary */
	unsigned char* ptr = (unsigned char*)filecontent;
	unsigned char* eptr = (unsigned char*)& filecontent[filelength];
	for (; ptr < eptr; cidx++, ptr++) {
		searchptr = &treehead;
		wordbufptr = wordbuf;
		while (searchptr->child[*ptr] != NULL) /* searchptr -> the ancestor */ {
			searchptr = searchptr->child[*ptr];
			*wordbufptr++ = *ptr++;
		}
		*wordbufptr++ = *ptr;

		/* maintain the tree structure */
		TreeNode* newnode = (TreeNode*)malloc(sizeof(TreeNode));
		searchptr->child[*ptr] = initnode(newnode, cidx, *ptr);

		/* insert the dictionary */
		insertdict(dict, cidx, wordbufptr - wordbuf, wordbuf, searchptr->idx);
		if ((ptr - (unsigned char*)filecontent) % (filelength / 100) == 0)
			progress_bar((ptr - (unsigned char *)filecontent) / (double)filelength, "Encode");
	}
	progress_bar((ptr - (unsigned char*)filecontent) / (double)filelength, "Encode");
	putchar('\n');
	QueryPerformanceCounter(&t2);
	printf("Encode Used :%fms\n", (t2.QuadPart - t1.QuadPart) * 1000.0 / tc.QuadPart);
	printf("Encode Speed :%fms / KB\n", (t2.QuadPart - t1.QuadPart) * 1000.0 / tc.QuadPart / filelength * 8 * 1024);
	
	/* if you want to show the dictionary, uncomment next line */
	// printsymtab(dict, cidx);

	/* Step 4 - Prepare encode buf and write encoded bits into it */
	/* Caution that the encodebuf may be larger than the origin filecontent buf */
	unsigned char* encodebuf = (unsigned char*)malloc(sizeof(unsigned char) * filelength * 8 * 2);
	BitStream* bs = bsnew(encodebuf, filelength);

	size_t written = 0;
	for (i = 1; i < cidx; i++) {
		size_t jlen = i == 1 ? 1 : (size_t)ceil(log2(i));
		int j;
		for (j = jlen - 1; j >= 0; j--)
			bswriteb(bs, (unsigned char)(dict[i].prev_idx >> j));
		for (j = 8 - 1; j >= 0; j--)
			bswriteb(bs, dict[i].append >> j);
		written += 8 + jlen;
	}
	size_t encodedlen = (written + 7) / 8;
	printf("Encoded File Size: %lu Btyes\n", (unsigned long)encodedlen);
	printf("Compress Rate: %lf%%\n", encodedlen / (double)filelength * 100);
	/* write back encode file */
	err = fopen_s(&fp, encodedfilename, "wb");
	fwrite(encodebuf, sizeof(char), encodedlen, fp);
	fclose(fp);
	printf("Encoded file %s written successfully\n", encodedfilename);

	/* Step 5 - Prepare decode buf and decode while generating the dictionary */
	bs->offset = 0;
	bs->ptr = bs->base;
	unsigned char* decodebuf = (unsigned char*)malloc(sizeof(unsigned char) * filelength);
	unsigned char* dptr = decodebuf;
	Entry* decodedict = (Entry*)malloc(sizeof(Entry) * filelength);
	decodedict[0].len = 0;
	size_t readb = 0;
	QueryPerformanceCounter(&t1);
	for (i = 1; readb < written; i++) {
		size_t jlen = i == 1 ? 1 : (size_t)ceil(log2(i));
		int j;
		size_t prev_idx = 0;
		unsigned char append = 0;
		for (j = jlen - 1; j >= 0; j--)
			prev_idx += bsgetb(bs) ? 1 << j : 0;
		for (j = 8 - 1; j >= 0; j--)
			append += (unsigned char)bsgetb(bs) << j;
		size_t k;
		for (k = 0; k < decodedict[prev_idx].len; k++)
			wordbuf[k] = decodedict[prev_idx].word[k];
		wordbuf[k++] = append;

		/* insert the dictionary */
		insertdict(decodedict, i, k, wordbuf, prev_idx);
		readb += 8 + jlen;
		
		if (i % (filelength / 100) == 0)
			progress_bar(readb / (double)written, "Decode");
	}
	progress_bar(readb / (double)written, "Decode"); putchar('\n');
	cidx = i;
	QueryPerformanceCounter(&t2);
	printf("Decode Used :%fms\n", (t2.QuadPart - t1.QuadPart) * 1000.0 / tc.QuadPart);
	printf("Decode Speed :%fms / KB\n", (t2.QuadPart - t1.QuadPart) * 1000.0 / tc.QuadPart / filelength * 8 * 1024);


	/* Step 6 - write each entry of dictionary into the decoded file */
	for (i = 1; i < cidx; i++) {
		size_t j;
		for (j = 0; j < decodedict[i].len; j++) {
			*dptr++ = decodedict[i].word[j];
		}
	}


	/* Step 7 - write back the encoded file */
	err = fopen_s(&fp, decodedfilename, "wb");
	fwrite(decodebuf, sizeof(char), filelength, fp);
	fclose(fp);
	printf("Decoded file %s written successfully\n", decodedfilename);

	/* Final Step - Clean up the memory */
	free(filecontent);
	free(dict);
	return 0;
}
